#include "ppdef.h"

module dineof_utils

contains

 !---------------------------------------------------------
 !----- subroutine valsvd ---------------------------------
 !---------------------------------------------------------

 ! computes the missing data values from the EOF basis     
 ! 
 !---------------------------------------------------------

subroutine valsvd(U,S,V,X,IEX,JEX,VAL,NN,IMISST) 

implicit none
 integer,intent(in)    :: IMISST
 real, intent(in)      :: U(:,:),V(:,:),S(:,:)
 real, intent(inout)   :: X(:,:)
 integer,intent(in)    :: NN,IEX(:),JEX(:)
 real,intent(out)      :: VAL
 integer               :: K,t


do t=1,IMISST
  VAL=0 
  DO K=1,NN 

    VAL=VAL+U(IEX(t),K)*S(K,1)*V(JEX(t),K) 
    
  ENDDO
  X(IEX(t),JEX(t))=VAL
enddo
 

end subroutine


 !cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
 subroutine dindiff(x,B,alpha,numit)
  use ufileformat
  implicit none

  real,intent(inout)         :: B(:)
  real,intent(in)            :: x(:)
  real, pointer              :: BF(:),xe(:)
  integer                    :: nsize,k,numit
  real                       :: alpha,valex


  nsize = size(B,1)

  allocate(BF(size(B,1)+1))

  allocate(xe(size(x,1)+1))
  BF=0

  ! extended x
  xe(1) = 1.5*x(1) - .5 * x(2)
  xe(2:nsize) = (x(1:nsize-1) + x(2:nsize))/2
  xe(nsize+1) = 1.5*x(nsize) - .5 * x(nsize-1)


  do k=1,numit
    !  F(2:nsize) = alpha * (f(2:nsize) - f(1:nsize-1));
    !  f = f +  (F(2:nsize+1) - F(1:nsize));  
    BF(2:nsize) = alpha * (B(2:nsize) - B(1:nsize-1))/(x(2:nsize) - x(1:nsize-1))
    B = B + (BF(2:nsize+1) - BF(1:nsize))/ (xe(2:nsize+1) - xe(1:nsize))
  end do

  deallocate(BF,xe)

 end subroutine dindiff

 !cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc


 !---------------------------------------------------------
 !----- subroutine writeMatrix ----------------------------
 !---------------------------------------------------------

 ! reshapes matrix into 3D file, then writes results  
 ! into separate files
 ! 
 !---------------------------------------------------------

 subroutine writeMatrix(X,xmean,valex,resultfnames,maskfile,norma,imax,jmax,first,DirOutput,fileMean,fileStd)
  use ufileformat
  use initfile
  implicit none

  real,intent(in)                  :: X(:,:)
  real,intent(in),optional         :: fileMean(:),fileStd(:)
  integer,intent(in)               :: imax(:),jmax(:),first(:)
  real,intent(in)                  :: xmean
  real,intent(inout)               :: valex
  integer,intent(in)               :: norma

  character(len=200),intent(in)    :: resultfnames(:),maskfile(:)
  character(len=200),intent(in)    :: DirOutput

  real, allocatable                :: sst(:,:,:)
  real,pointer                     :: mask(:,:)
  real, parameter                  :: valexc = 9999.
  real                             :: mean,stddev,var
  integer                          :: i,j,s,t,q,nbvars,NN,N
  character(len=200)               :: initfilename


  N = size(X,2)

  nbvars=size(resultfnames)
  call getarg(1,initfilename)

  !Reshape the matrix into a 3D matrix (two spatial dimensions and one temporal dimension)
  !---------------------------------------------------------------------------------------
  do q=1,nbvars

    allocate(sst(imax(q),jmax(q),N))

    if( presentInitValue(initfilename,'mask')) then
      call uload(maskfile(q),mask,valex)
    else
      allocate(mask(size(sst,1),size(sst,2)))
      mask = mask+1 
    endif

      where (mask.eq.0) mask = valexc

      do t=1,N

      s = first(q)

      do i=1,imax(q)
        do j=1,jmax(q)

          if(mask(i,j).ne.valexc) then

            sst(i,j,t)=X(s,t);

            if(norma.eq.1) then
              sst(i,j,t)=(sst(i,j,t)*fileStd(q))+fileMean(q)
            end if

            s = s+1
          else
            sst(i,j,t)=valexc;
          endif
        enddo
      enddo
    enddo



    NN = count(sst.ne.valexc)
    mean = sum(sst,sst.ne.valexc)/NN
    var = sum((sst-mean)**2,sst.ne.valexc)/NN
    stddev = sqrt(var)


    if(norma.eq.1) then
      !   -----------------------------------------
      !   some statistics about the filled matrices
      !   -----------------------------------------

      write(stdout,*)
      write(stdout,*)'mean ',mean
      write(stdout,*)'Standard deviation',stddev,fileStd,fileMean
      write(stdout,*)
    end if

    call usave(trim(resultfnames(q)),sst(:,:,:),valexc)

    deallocate(sst,mask)

  enddo

 end subroutine writeMatrix



 !---------------------------------------------------------
 !----- subroutine writeSVD ----------------------------
 !---------------------------------------------------------

 ! writes EOFs
 ! 
 !---------------------------------------------------------


 subroutine writeSVD(DirOutput,eofvfname,eofsfname,s,v,valc,sumVarEnd)
  use ufileformat
  use netcdf
  implicit none

  character(len=200),intent(in)     :: DirOutput,eofvfname,eofsfname
  real,intent(in)                   :: s(:),v(:,:),valc(:),sumVarEnd

  integer                           :: i,j,k,ncid,status
  real, parameter                   :: valex = 9999.
  real, parameter                   :: valexc = 9999.
  real                              :: varex(size(s))


  ! call usave(trim(resultfnames(q)),sst(:,:,:),valexc)

  status = nf90_create(path = trim(DirOutput)//'/DINEOF_diagnostics.nc', cmode = nf90_clobber, ncid = ncid)
  if (status /= nf90_noerr) call handle_err(status)

  status = nf90_close(ncid)
  if (status /= nf90_noerr) call handle_err(status)

  !   Variance detailed
  !   _______________

  varex= 100.0*s**2/sumVarEnd
write(*,*)'sumvarEnd ',sumVarEnd
  call usave(trim(DirOutput)//'/DINEOF_diagnostics.nc#varEx',varex,valex)

  write(stdout,*) 'Sum of the squares of the singular values of the ',size(s), 'EOFs retained', sum(varex)

  !
  !   Singular values
  !   _______________


  call usave(trim(DirOutput)//'/DINEOF_diagnostics.nc#vlsng',s,valex)

 
  !
  !   Temporal modes 
  !   _______________

  call usave(eofvfname,v,valex)

  !---------------------------------------------------------------------------------------

  !   Valc: write expected error for each mode
  !   ----------------
  call usave(trim(DirOutput)//'/DINEOF_diagnostics.nc#valc',valc,valexc)   
  !---------------------------------------------------------------------------------------


  !stop      

500 format(a72)
501 format(e10.4)
502 format(3e10.4)
503 format(300(e10.4,1X))
504 format(a5,i3,a3,f8.4,a2)

 end subroutine writeSVD

subroutine clear_files(filenames)
 use netcdf
 character(len=200),intent(in)    :: filenames
 integer                          :: status,ncid



  status = nf90_create(path = trim(filenames), cmode = nf90_clobber, ncid = ncid)
  if (status /= nf90_noerr) call handle_err(status)
  
  status = nf90_close(ncid)
  if (status /= nf90_noerr) call handle_err(status)

 
 end subroutine clear_files



 subroutine handle_err(status)
 use netcdf
        integer, intent ( in) :: status
     
        if(status /= nf90_noerr) then
          print *, trim(nf90_strerror(status))
          stop "Stopped"
        end if
 end subroutine handle_err


 !---------------------------------------------------------
 !----- subroutine dineof_merging ----------------------------
 !---------------------------------------------------------

 ! analysis of two or more data sets using EOF basis from DINEOF
 ! 
 !---------------------------------------------------------


 subroutine dineof_merging(X,IEX,JEX,U,S,fileStd,fileMean,data_lon,data_lat,data_stderr,mask,data2,lon2,lat2,time2,std2)
  use ufileformat
  use initfile
  use ndgrid
  use matoper
  use rrsqrt
  implicit none
  
  ! parameters
  real,intent(inout)            :: X(:,:),mask(:,:)
  real,intent(in)               :: U(:,:),fileStd(:),fileMean(:)
  real,intent(in)               :: S(:)
  real,intent(in)               :: data_lon(:,:),data_lat(:,:),data_stderr(:,:,:)
  real,intent(in)               :: data2(:),lon2(:),lat2(:),time2(:)
  real,intent(inout)            :: std2(:)

  integer,intent(in)            :: IEX(:),JEX(:)
 
!  character(len=200),intent(in) :: initfilename

  ! local variables
!  character(len=200),pointer    :: addfile(:),addfilelat(:),addfilelon(:),addfiletime(:)
!  character(len=200),pointer    :: addfilestderr(:),filelat(:),filelon(:),filestderr(:)

  ! additional data (value, lon, lat, time and error standard deviation)
!  real, pointer                 :: data2(:),lon2(:),lat2(:),time2(:),std2(:)
  real, pointer                 :: X_grid(:,:,:),U_grid(:,:,:)
  real, allocatable             :: data_stderr_packed(:,:)
  real                          :: valex, tmp
  real, parameter               :: valexc = -9999.
  real, allocatable             :: sst2(:),yo(:),HSf(:,:),invsqrtR(:),xa_xf(:,:),Sf(:,:),Hxf(:)
  integer                       :: i,j,k,p,nobs,t
  logical :: out
  type(grid) :: g


 !read file names
!   call getInitValue(initfilename,'add_lon',addfilelon)
!   call getInitValue(initfilename,'add_lat',addfilelat)
!   call getInitValue(initfilename,'add_time',addfiletime)
!   call getInitValue(initfilename,'add_data',addfile)
!   call getInitValue(initfilename,'add_std',addfilestderr)

!   call getInitValue(initfilename,'data_lon',filelon)
!   call getInitValue(initfilename,'data_lat',filelat)
!   call getInitValue(initfilename,'data_std',filestderr)

 !load data
 !considering 1 extra data file. Modify for more datafiles later
 
!   call uload(addfilelon(1),lon2,valex)
!   call uload(addfilelat(1),lat2,valex)
! assume time2 are time indexes
!   call uload(addfiletime(1),time2,valex)
!   call uload(addfile(1),data2,valex)
!   call uload(addfilestderr(1),std2,valex)

!   call uload(filelon(1),data_lon,valex)
!   call uload(filelat(1),data_lat,valex)
!   call uload(filestderr(1),data_stderr,valex)

!   if( presentInitValue(initfilename,'mask')) then
!     call uload(maskfile(1),mask,valex)
!     where (mask.eq.0) mask = valexc
!   else
     ! check!!!
!     allocate(mask(size(X_grid,1),size(X_grid,2)))
!     mask = mask+1 
!   endif
   
   ! normalize additional data
!   data2 = (data2 - fileMean(1)) / fileStd(1)
  
   allocate(data_stderr_packed(size(X,1),size(X,2)))
   data_stderr_packed = dineof_pack(data_stderr,mask /= valexc)
 
 ! unpack X and U
 !  write(6,*) 'range X ',minval(X),maxval(X)
   allocate(X_grid(size(mask,1),size(mask,2),size(X,2)))

   !call dineof_unpack2(X,mask /= valexc,X_grid,fileStd,fileMean,1)
   
   X_grid = dineof_unpack(X,mask /= valexc,valexc)

   call dineof_unpack2(U,mask /= valexc,U_grid,(/ 1./),(/0./),1)
!   write(6,*) 'range sst ',minval(X_grid),maxval(X_grid),fileStd,count(mask == valexc)

!   write(6,*) 'check pack ',maxval(abs(U - dineof_pack(U_grid,mask /= valexc)))
!   write(6,*) 'check pack ',maxval(abs(U_grid - dineof_unpack(U,mask /= valexc,9999.)))

   ! put valexc at the location of missing values ("clouds" and cross-validation points)
   do i=1,size(iex)
     X(IEX(i),JEX(i)) = valexc
   end do
   
! write(*,*)'after unpack'
!   write(*,*)'size matrices ',shape(X),shape(U)

   allocate(sst2(size(data2)))

!   write(6,*) 'shape ',shape(X_grid),shape(data_lon),shape(data_lat)

   ! initialize spatial grid (data_lon,data_lat) for interpolation

   call initgrid(g,data_lon,data_lat,mask == valexc)
!   call usave('test_X_grid.dat',X_grid,-9999.)
!   call usave('test_X.dat',X,-9999.)
!   call usave('test_iex.dat',real(iex),-9999.)
!   call usave('test_jex.dat',real(jex),-9999.)

   ! interpolate SST on lon2,lat2,time2 -> sst2
   do i=1,size(data2)
     call interp(g,X_grid(:,:,int(time2(i))),(/ lon2(i),lat2(i) /),sst2(i),out)

!     write(6,*) 'adddata lon lat ',lon2(i),lat2(i),time2(i),sst2(i),out

     ! not in grid then, give weight = 0
     if (out) then
       std2(i) = 1e10;
       sst2(i) = 0;
     end if
     
   end do

   ! since X are anomalies, the error covariance P is 
   ! Pf = 1/(n-1) X * X^T
   ! for rrsqrt we need the square root matrix S of the error covariance Pf = Sf * Sf^T
   ! SVD decomposition of X:
   ! X = U Sigma V^T
   ! Thus Sf is:
   ! Sf = 1/sqrt(n-1) U Sigma
   

   allocate(xa_xf(size(X,1),size(X,2)),Sf(size(U,1),size(U,2)))

!   call usave('test_U.dat',U_grid,9999.)
!   call usave('test_S.dat',S,9999.)

   ! the columns of Sf are the columns of U time the corresponding singular value
   Sf = 1./sqrt(size(X,2)-1.) * (U .xd. S);

!   call usave('test_stddevxf.dat', dineof_unpack(reshape(sum(Sf**2,2),[size(X,1),1]),mask /= valexc,9999.) ,9999.)

   ! loop over time
   do j = 1,size(X_grid,3)
   !do j = 1,1
     ! total number of observation for ith time 
     nobs = count(time2 == j) + size(X,1) - count(JEX == j)
!     write(6,*) 'nobs ',nobs,count(time2 == j) ,size(X,1) ,count(JEX == j)

!     write(6,*) 'test ',X(IEX,j)
!     write(6,*) 'test ',count(JEX == j),count(X(:,j) == valexc)
!     write(6,*) 'data_stderr ',shape(data_stderr)

     ! d: vector with all observation
     ! HSf: observed part of EOFs
     allocate(yo(nobs),HSf(nobs,size(U,2)),invsqrtR(nobs),Hxf(nobs))

     Hxf = 0

     ! collect all data that is present in j-th time instance
     k = 0
     do i = 1,size(X,1)
       if (X(i,j) /= valexc) then
         k = k+1
         yo(k) = X(i,j)
         HSf(k,:) = Sf(i,:)
         invsqrtR(k) = 1./sqrt(data_stderr_packed(i,j))
       end if             
     end do

     ! loop over all additional (e.g. in-situ) data
     do i = 1,size(time2)
       if (time2(i) == j) then
         k = k+1
         yo(k) = data2(i)

         do p=1,size(HSf,2)
           call interp(g,U_grid(:,:,p),(/ lon2(i),lat2(i) /),HSf(k,p),out)
         end do

         ! check out
         if (.not.out) then

           invsqrtR(k) = 1./sqrt(std2(i))

  !         write(6,*) 'check val ',yo(k),HSf(k,:)
         else
 !          write(6,*) 'ignore data point',lon2(i),lat2(i)
           ! change R                           
           invsqrtR(k) = 0.
         end if
       end if
     end do

!     write(6,*) 'should be the same ',k,nobs

! Hxf background estimate at obs. location (here = 0 everywhere)
! yo all observations from X and data2
! Sf: EOF * singular value
! HSf: Sf at obs. location
 
!    call usave('test_invsqrtR.dat',invsqrtR ,-9999.)
 !    call usave('test_yo.dat',yo ,-9999.)
 !    call usave('test_HSf.dat',HSf ,-9999.)
 !    call usave('test_Hxf.dat',Hxf ,-9999.)
 
     call analysisIncrement(Hxf,yo,Sf,HSf,invsqrtR, xa_xf(:,j))
 !    write(6,*) 'analysisIncrement ',minval(xa_xf(:,j)),maxval(xa_xf(:,j))
 !    write(6,*) 'analysisIncrement ',minval(yo),maxval(yo)
 !    write(6,*) 'analysisIncrement ',minval(invsqrtR),maxval(invsqrtR)
 !    write(6,*) 'analysisIncrement HSf',minval(HSf),maxval(HSf)

     
     
     deallocate(yo,HSf,invsqrtR,Hxf)
   end do

   do t=1,size(IEX,1)
     X(IEX(t),JEX(t)) = xa_xf(IEX(t),JEX(t))
    end do

!   call usave('test_xa.dat',dineof_unpack(xa_xf,mask /= valexc,9999.) ,9999.)
!   call usave('test_Sf.dat',dineof_unpack(Sf,mask /= valexc,9999.) ,9999.)
!   call usave('test_Sfp.dat',Sf ,9999.)

   deallocate(X_grid,U_grid)
 end subroutine dineof_merging   


!---------------------------------------------------------
!----- subroutine dineof_unpack2 ----------------------------
!---------------------------------------------------------

 ! Reshapes from 2D field to 3D field
 ! 
 ! done for 1 matrix for the moment: generalize to multivariate case
 !---------------------------------------------------------

 subroutine dineof_unpack2(X,mask,field,fileStd,fileMean,norma)
  implicit none

  real,intent(in)                   :: X(:,:),fileStd(:),fileMean(:)
  logical,intent(in)                :: mask(:,:) ! .true. = sea (valid) and .false. land (no-valid)
  integer                           :: norma

  real,pointer                      :: field(:,:,:)
  integer                           :: i,j,s,t,N
  real,parameter                    :: valexc = 9999.

  

  N = size(X,2)
  allocate(field(size(mask,1),size(mask,2),N))

  do t = 1,N    
    s = 1
    do i = 1,size(mask,1)
      do j = 1,size(mask,2)

        if (mask(i,j)) then          
          field(i,j,t) = X(s,t)

          if(norma.eq.1) then
            field(i,j,t) = (field(i,j,t)*fileStd(1))+fileMean(1)
          end if

          s = s+1
        else
          field(i,j,t) = valexc
        endif
      enddo
    enddo
  enddo

 end subroutine dineof_unpack2


 function dineof_pack(field,mask) result(X)
  implicit none
  real,intent(in)                   :: field(:,:,:)
  logical,intent(in)                :: mask(:,:)
  real                  :: X(count(mask),size(field,3))

  integer                           :: i,j,s,t

  do j=1,size(X,2)
    ! to do:
    ! need to change convention to avoid transpose here

    X(:,j) = pack(transpose(field(:,:,j)),transpose(mask))
  end do


  ! do t = 1,size(X,2)
  !   s = 1
  !   do i = 1,size(mask,1)
  !     do j = 1,size(mask,2)

  !       if (mask(i,j)) then          
  !         X(s,t) = field(i,j,t)
  !         s = s+1
  !       end if
  !     end do
  !   end do
  ! end do

 end function dineof_pack


! expand data in X using land-sea mask

 function dineof_unpack(X,mask,fillvalue) result(field)
  implicit none
  logical,intent(in)    :: mask(:,:)
  real                  :: X(:,:)
  real                  :: field(size(mask,1),size(mask,2),size(X,2))
  real                  :: fillvalue
  integer                           :: i,j,s,t

  field = fillvalue
  do j=1,size(X,2)
    ! to do:
    ! need to change convention to avoid transpose here

    field(:,:,j) = transpose(unpack(X(:,j),transpose(mask),transpose(field(:,:,j))))
  end do


  ! do t = 1,size(X,2)
  !   s = 1
  !   do i = 1,size(mask,1)
  !     do j = 1,size(mask,2)

  !       if (mask(i,j)) then          
  !         X(s,t) = field(i,j,t)
  !         s = s+1
  !       end if
  !     end do
  !   end do
  ! end do

 end function dineof_unpack

end module dineof_utils
